package test.demo;

import cn.hutool.core.util.ArrayUtil;
import tech.waterism.modelbase.PredictBase;
import tech.waterism.modelbase.PredictResult;
import com.alibaba.fastjson.JSON;

import java.util.*;

/**
 * 马斯京根河道连续演算法
 *
 * @author 初京刚
 * @version 1.0
 */
public class MSK extends PredictBase {

    /**
     * 马法初始流量，河段下断面初始流量
     */
    private double[] QXInitial;

    /**
     * 子河段洪水波传播时间
     */
    private double KE;

    /**
     * 河段数
     */
    private int N;

    /**
     * 子河段流量比重因子
     */
    private double XE;

    /**
     * 降水时间
     */
    private Date[] floodTm;

    /**
     * 预报开始时间
     */
    private Date startTm;

    /**
     * 预报开始数组下标
     */
    private int startIndex;

    /**
     * 降水序列长度
     */
    private Integer floodRainRange;

    /**
     * 降水时段长
     */
    private double timeInterval;

    /**
     * 上游断面流量过程
     */
    private double[] Q;

    /**
     * 马法初始流量，河段下断面各时段初始流量
     */
    private double[][] QX0;

    public MSK(Map<String, String> data) {
        super();
        QXInitial = JSON.parseObject(data.get("QX"), double[].class);

        XE = Double.parseDouble(data.get("X"));
        KE = Double.parseDouble(data.get("KK"));
        N = Integer.parseInt(data.get("MP"));

        floodTm = JSON.parseObject(data.get("dt"), Date[].class);   // input[i].data.dt
        startTm  = JSON.parseObject(data.get("start"), Date.class);
        startIndex = ArrayUtil.indexOf(floodTm, startTm);
        Q = JSON.parseObject(data.get("q"), double[].class);      // input[i].data.q
        timeInterval = Double.parseDouble(data.get("clen"));        // index.clen
        floodRainRange = Q.length;
        // 修正QXInitial
        if(QXInitial.length > 0 && Q.length > 0 && startIndex < Q.length){
            double rate = 1.0;
            double difference = 0.0;
            boolean doPlus = false;
            if(QXInitial[0] != 0){
                rate = Q[startIndex] / QXInitial[0];
            }else{
                difference = Q[startIndex];
                doPlus = true;
            }
            for(int i = 0; i < QXInitial.length; i++){
                if(doPlus){
                    QXInitial[i] += difference;
                }else{
                    QXInitial[i] *= rate;
                }
            }
        }
    }

    /**
     * 模型计算
     */
    public Map<String, double[]> MSKmodel() {

        // 河道演进
        double[] QSim = muskingumSeg(Q);

        Map<String, double[]> simulationResult = new HashMap<>();
        simulationResult.put("Q", Q);
        simulationResult.put("QSim", QSim);
        return simulationResult;
    }

    /**
     * 河道演进-分段马斯京根
     *
     * @param Qr 河网汇流-单元面积出口流量（m3/s）
     * @return 河道演进-单元面积出口流量（m3/s）
     */
    private double[] muskingumSeg(double[] Qr) {
        //        double[] QC = new double[Qr.length];
        double[] QC = new double[N];//河段下断面时段初始流量
        double[] QRouting = new double[Qr.length];//河道演进结果
        QX0 = new double[Qr.length + 1][N];// 新增逻辑:存储每一步的初值
        for (Integer j = 0; j < N; j++) {
            if (j < QXInitial.length) {
                QC[j] = QXInitial[j];
            } else {
                QC[j] = QC[j - 1];
            }
            QX0[0][j] = QC[j];
        }

        double CD = 0.5 * timeInterval + KE - KE * XE;
        double C0 = (0.5 * timeInterval - KE * XE) / CD;
        double C1 = (0.5 * timeInterval + KE * XE) / CD;
        double C2 = 1 - C0 - C1;

        double QI1 = 0, QI2 = 0, QO1 = 0, QO2 = 0;//河段上、下断面时段初始流量和时段末流量
        if (C0 >= 0.0 && C2 >= 0.0) {// C0和C2需大于等于0，才满足马斯京根法的适用条件
            for (Integer i = startIndex; i < Qr.length; i++) {
                if (i == startIndex) {
                    QI1 = Qr[i];
                } else {
                    QI1 = Qr[i - 1];
                }
                QI2 = Qr[i];

                // 分段计算
                if (N > 0) {
                    for (Integer j = 0; j < N; j++) {
                        QX0[i][j] = QC[j];// 新增逻辑:存储每一步的初值
                        QO1 = QC[j];
                        QO2 = C0 * QI2 + C1 * QI1 + C2 * QO1;
                        QI1 = QO1;
                        QI2 = QO2;
                        QC[j] = QO2;
                    }
                } else {

                    // 新增逻辑:存储每一步的初值
                    QX0[i] = QXInitial;

                    QO2 = Qr[i];
                }

                if (QO2 < 0.0001f)
                    QO2 = 0.0f;
                QRouting[i] = QO2;
            }

            for (Integer j = 0; j < N; j++) {
                QX0[Qr.length][j] = QC[j];
            }
        } else {// 否则，就不考虑河道演进了
            System.arraycopy(Qr, 0, QRouting, 0, Qr.length);
        }

        return QRouting;
    }

    private List<String> toStringList(double[][] array0) {
        String[] array = new String[array0.length];
        for (int i = 0; i < array.length; i++) {
            array[i] = Arrays.toString(array0[i]);
        }
        return Arrays.asList(array);
    }

    @Override
    public void init() {

    }

    @Override
    public PredictResult predict() {
        PredictResult result = new PredictResult();
        Map<String, double[]> map = MSKmodel();

        result.setQSim(map.get("QSim"));

        double[] RSim = map.get("QSim");

        for (int i = 0; i < RSim.length; i++) {
            System.out.println(RSim[i]);
        }

        // 更新状态
        List<String> QX = toStringList(QX0);
        Map<String, List<String>> status = new HashMap<>();
        status.put("QX", QX);
        result.setNewStatus(status);
        return result;
    }
}